// SPDX-License-Identifier: Apache-2.0
//
// Copyright © 2017 Trust Wallet.

use crate::SOLANA_ALPHABET;
use serde::de::Error as DeError;
use serde::{Deserialize, Deserializer, Serialize, Serializer};
use std::fmt;
use std::str::FromStr;
use tw_coin_entry::coin_entry::CoinAddress;
use tw_coin_entry::error::prelude::*;
use tw_encoding::base58;
use tw_hash::{as_byte_sequence, sha2, H256};
use tw_keypair::{ed25519, tw};
use tw_memory::Data;

pub const MAX_SEEDS: usize = 16;
pub const MAX_SEED_LEN: usize = H256::LEN;
const PDA_MARKER: &[u8; 21] = b"ProgramDerivedAddress";

#[derive(Clone, Copy, Default, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub struct SolanaAddress {
    bytes: H256,
}

impl SolanaAddress {
    pub fn with_public_key(public_key: &tw::PublicKey) -> AddressResult<SolanaAddress> {
        let bytes = public_key
            .to_ed25519()
            .ok_or(AddressError::PublicKeyTypeMismatch)?
            .to_bytes();
        Ok(SolanaAddress { bytes })
    }

    pub fn with_public_key_ed25519(public_key: &ed25519::sha512::PublicKey) -> SolanaAddress {
        SolanaAddress {
            bytes: public_key.to_bytes(),
        }
    }

    pub fn with_public_key_bytes(bytes: H256) -> SolanaAddress {
        SolanaAddress { bytes }
    }

    pub fn bytes(&self) -> H256 {
        self.bytes
    }

    /// Find a valid [program derived address][pda] and its corresponding bump seed.
    ///
    /// [pda]: https://solana.com/docs/core/cpi#program-derived-addresses
    ///
    /// Program derived addresses (PDAs) are account keys that only the program,
    /// `program_id`, has the authority to sign. The address is of the same form
    /// as a Solana `Pubkey`, except they are ensured to not be on the ed25519
    /// curve and thus have no associated private key. When performing
    /// cross-program invocations the program can "sign" for the key by calling
    /// [`invoke_signed`] and passing the same seeds used to generate the
    /// address, along with the calculated _bump seed_, which this function
    /// returns as the second tuple element. The runtime will verify that the
    /// program associated with this address is the caller and thus authorized
    /// to be the signer.
    pub fn find_program_address(
        seeds: &[&[u8]],
        program_id: SolanaAddress,
    ) -> Option<SolanaAddress> {
        let mut bump_seed = [u8::MAX];
        for _ in 0..u8::MAX {
            let mut seeds_with_bump = seeds.to_vec();
            seeds_with_bump.push(&bump_seed);
            match Self::create_program_address(&seeds_with_bump, program_id) {
                Ok(Some(address)) => return Some(address),
                // Try to re-compute the program address with a different seed.
                Ok(None) => (),
                Err(_) => return None,
            }
            // Try to re-compute the program address with a different seed.
            bump_seed[0] -= 1;
        }
        None
    }

    /// Create a valid [program derived address][pda] without searching for a bump seed.
    ///
    /// [pda]: https://solana.com/docs/core/cpi#program-derived-addresses
    ///
    /// Because this function does not create a bump seed, it may unpredictably
    /// return an error for any given set of seeds and is not generally suitable
    /// for creating program derived addresses.
    ///
    /// However, it can be used for efficiently verifying that a set of seeds plus
    /// bump seed generated by [`find_program_address`] derives a particular
    /// address as expected. See the example for details.
    ///
    /// See the documentation for [`find_program_address`] for a full description
    /// of program derived addresses and bump seeds.
    ///
    /// [`find_program_address`]: Pubkey::find_program_address
    pub fn create_program_address(
        seeds: &[&[u8]],
        program_id: SolanaAddress,
    ) -> AddressResult<Option<SolanaAddress>> {
        if seeds.len() > MAX_SEEDS {
            return Err(AddressError::Internal);
        }
        if seeds.iter().any(|seed| seed.len() > MAX_SEED_LEN) {
            return Err(AddressError::Internal);
        }

        let mut data_to_hash = Vec::new();

        // concatenate seeds
        for seed in seeds {
            data_to_hash.extend_from_slice(seed);
        }
        // Append `program_id`.
        data_to_hash.extend_from_slice(program_id.bytes.as_slice());
        data_to_hash.extend_from_slice(PDA_MARKER);

        let hash = H256::try_from(sha2::sha256(&data_to_hash).as_slice())
            .expect("sha256 must return 32 bytes");

        // The given hash (aka new public key) must not be on the ed25519 elliptic curve.
        match ed25519::sha512::PublicKey::try_from(hash.as_slice()) {
            Ok(_) => Ok(None),
            Err(_) => Ok(Some(SolanaAddress::with_public_key_bytes(hash))),
        }
    }
}

impl CoinAddress for SolanaAddress {
    #[inline]
    fn data(&self) -> Data {
        self.bytes.to_vec()
    }
}

impl FromStr for SolanaAddress {
    type Err = AddressError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        let bytes =
            base58::decode(s, SOLANA_ALPHABET).map_err(|_| AddressError::FromBase58Error)?;
        let bytes = H256::try_from(bytes.as_slice()).map_err(|_| AddressError::InvalidInput)?;
        Ok(SolanaAddress { bytes })
    }
}

impl From<&'static str> for SolanaAddress {
    fn from(s: &'static str) -> Self {
        SolanaAddress::from_str(s).unwrap()
    }
}

impl fmt::Debug for SolanaAddress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}", self)
    }
}

impl fmt::Display for SolanaAddress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let encoded = base58::encode(self.bytes.as_slice(), SOLANA_ALPHABET);
        write!(f, "{}", encoded)
    }
}

impl Serialize for SolanaAddress {
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        if serializer.is_human_readable() {
            return self.to_string().serialize(serializer);
        }
        as_byte_sequence::serialize(&self.bytes(), serializer)
    }
}

impl<'de> Deserialize<'de> for SolanaAddress {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
        D: Deserializer<'de>,
    {
        if deserializer.is_human_readable() {
            let addr_str = String::deserialize(deserializer)?;
            return SolanaAddress::from_str(&addr_str)
                .map_err(|e| DeError::custom(format!("{e:?}")));
        }
        let bytes = as_byte_sequence::deserialize(deserializer)?;
        Ok(SolanaAddress { bytes })
    }
}
